#ifndef EXAMPLES_MATMULSUB_KERNEL_H
#define EXAMPLES_MATMULSUB_KERNEL_H
#include "kernel_operator.h"
#include "lib/matmul_intf.h"

using namespace matmul;

__aicore__ inline uint32_t Ceiling(uint32_t a, uint32_t b)
{
    return (a + b - 1) / b;
}
template<typename aType, typename bType, typename cType> class KernelMatMulSub {
public:
    __aicore__ inline KernelMatMulSub(){};
     __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR x3, GM_ADDR y, GM_ADDR workspace,
                                const TCubeTiling &tiling, int8_t broadcast, AscendC::TPipe *pipe);
    __aicore__ inline void Process(AscendC::TPipe *pipe);

    __aicore__ inline void MatmulCompute();
    __aicore__ inline void LeakyReluCompute();
    __aicore__ inline void CopyOut(uint32_t count);
    __aicore__ inline void CalcOffset(int32_t blockIdx, const TCubeTiling &tiling, int32_t &offsetA, int32_t &offsetB,
                                      int32_t &offsetC);

    Matmul<MatmulType<TPosition::GM, CubeFormat::ND, DTYPE_X1>, 
           MatmulType<TPosition::GM, CubeFormat::ND, DTYPE_X2>,
           MatmulType<TPosition::GM, CubeFormat::ND, DTYPE_Y>>
        matmulObj;
    
    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_X3> x3Gm;
    GlobalTensor<DTYPE_Y> yGm;
    GlobalTensor<DTYPE_Y> y1Gm;

    TCubeTiling tiling;

    TQue<QuePosition::VECIN, 2> inQueueX3;
    TQue<QuePosition::VECOUT, 2> outQueueY;

    int8_t broadcast;
    int tailM;
    int tailN;
};

template <typename aType, typename bType, typename cType>
__aicore__ inline void
KernelMatMulSub<aType, bType, cType>::Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR x3, GM_ADDR y, GM_ADDR workspace,
                                                       const TCubeTiling &tiling, int8_t broadcast, AscendC::TPipe *pipe)
{
     this->tiling = tiling;
    x1Gm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X1 *>(x1), tiling.M * tiling.Ka);
    x2Gm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X2 *>(x2), tiling.Kb * tiling.N);
    x3Gm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X3 *>(x3), tiling.M * tiling.N);
    yGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y), tiling.M * tiling.N);
    y1Gm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y));

    //196352
    pipe->InitBuffer(inQueueX3, 2, 11776 * sizeof(DTYPE_X3));
    pipe->InitBuffer(outQueueY, 2, 11776 * sizeof(DTYPE_Y)); // Init output buffer.
    

   this->broadcast =broadcast;
    
    int offsetA = 0;
    int offsetB = 0;
    int offsetC = 0;
    // int offsetBias = 0;
    CalcOffset(GetBlockIdx(), tiling, offsetA, offsetB, offsetC);//, offsetBias); // Calculate the gm offset based on the blockidx.
    // yGm1 = yGm;
    // printf("offset:%d %d %d\r\n", offsetA, offsetB, offsetC);
    // printf("k:%d %d %d %d\r\n", tiling.M, tiling.N, tiling.Ka, tiling.Kb);
    // printf("base:%d %d %d\r\n", tiling.baseM, tiling.baseM, tiling.baseM);
    // printf("singleCore:%d %d %d\r\n", tiling.singleCoreM, tiling.singleCoreM, tiling.singleCoreM);
    // printf("DEFAULT_C0_SIZE:%d\r\n", DEFAULT_C0_SIZE);

    x1Gm = x1Gm[offsetA];
    x2Gm = x2Gm[offsetB];
    yGm = yGm[offsetC];

    if (GetSysWorkSpacePtr() == nullptr) {
        // printf("error:GetSysWorkSpacePtr() == nullptr\r\n");
        return;
    }
    
}

template <typename aType, typename bType, typename cType>
__aicore__ inline void KernelMatMulSub<aType, bType, cType>::Process(AscendC::TPipe *pipe)
{
    uint32_t computeRound = 0;

    matmulObj.SetTensorA(x1Gm);
    matmulObj.SetTensorB(x2Gm);
    // matmulObj.SetBias(biasGlobal);
    matmulObj.ClearBias(); 
    matmulObj.IterateAll<true>(yGm);
    // while (matmulObj.Iterate<true>()) {
    //     MatmulCompute();
    //     LeakyReluCompute();
    //     CopyOut(computeRound);
    //     computeRound++;
    // }
    matmulObj.End();

    // DataSyncBarrier<MemDsbT::ALL>();
    // printf("2");
    // SyncAll();
    SetAtomicAdd<DTYPE_Y>();
    
    if(broadcast == 0)
    {
        SyncAll();
        // if(GetBlockIdx() == 0)
        {
            uint32_t loop =  (tiling.M * tiling.N) / 11776;
            uint32_t wei = (tiling.M * tiling.N) % 11776;
            uint32_t wei_32 = (wei+7)/8*8;

            uint32_t loop_small_num = loop / GetBlockNum()/2;
            uint32_t loop_big = loop % (GetBlockNum()*2);
            LocalTensor<DTYPE_X3> x3Local;
            LocalTensor<DTYPE_Y> yLocal;

            uint32_t loop_start;
            uint32_t loop_end;
            if(GetBlockIdx() < loop_big)
            {
                loop_start = GetBlockIdx() * (loop_small_num + 1);
                loop_end = loop_start + loop_small_num + 1;
            }
            else
            {
                loop_start = GetBlockIdx() * (loop_small_num + 1) - (GetBlockIdx() - loop_big);
                loop_end = loop_start + loop_small_num;
                
            }
            // printf("%d %d %d %d %d\r\n",loop_start, loop_end, loop_big, loop, wei);
            for(int i=loop_start;i<loop_end;i++)
            {
                x3Local = inQueueX3.AllocTensor<DTYPE_X3>();
                DataCopy(x3Local, x3Gm[11776*i], 11776);
                inQueueX3.EnQue(x3Local);

                yLocal = outQueueY.AllocTensor<DTYPE_Y>();
                x3Local = inQueueX3.DeQue<DTYPE_X3>();
                // Duplicate(yLocal, (DTYPE_Y)0, 11776);
                // Sub(yLocal, yLocal, x3Local, 11776);
                Muls(yLocal, x3Local, (DTYPE_X3)-1, 11776);
                inQueueX3.FreeTensor(x3Local);
                outQueueY.EnQue(yLocal);

                yLocal = outQueueY.DeQue<DTYPE_Y>();
                DataCopy(y1Gm[11776*i], yLocal, 11776);
                outQueueY.FreeTensor(yLocal);
            }
            if(GetBlockIdx() == 0 && wei!=0)
            {
                x3Local = inQueueX3.AllocTensor<DTYPE_X3>();
                DataCopy(x3Local, x3Gm[11776*loop], wei_32);
                inQueueX3.EnQue(x3Local);

                yLocal = outQueueY.AllocTensor<DTYPE_Y>();
                x3Local = inQueueX3.DeQue<DTYPE_X3>();
                // Duplicate(yLocal, (DTYPE_Y)0, wei_32);
                // Sub(yLocal, yLocal, x3Local, wei_32);
                Muls(yLocal, x3Local, (DTYPE_X3)-1, wei_32);
                inQueueX3.FreeTensor(x3Local);
                outQueueY.EnQue(yLocal);

                yLocal = outQueueY.DeQue<DTYPE_Y>();
                DataCopy(y1Gm[11776*loop], yLocal, wei_32);
                // DataCopyExtParams copyParams {1, static_cast<uint32_t>(wei * sizeof(DTYPE_Y)), 0, 0, 0}; // 结构体DataCopyExtParams最后一个参数是rsv保留位
                // DataCopyPad(yGm[8192*loop], yLocal, copyParams);
                outQueueY.FreeTensor(yLocal);
            }
        }
    }
    else //if(std::is_same_v<DTYPE_X3, float>)
    {
        SyncAll();
        if(GetBlockIdx() == 0)
        {
            LocalTensor<DTYPE_X3> x3Local;
            LocalTensor<DTYPE_Y> yLocal;
            x3Local = inQueueX3.AllocTensor<DTYPE_X3>();
            yLocal = outQueueY.AllocTensor<DTYPE_Y>();
            uint32_t loop =  tiling.M * tiling.N;
            uint32_t source = tiling.N;
            for(uint32_t i=0;i<loop;i++)
            {
                x3Local.SetValue(0, x3Gm.GetValue(i%source));
                yLocal.SetValue(0, yGm.GetValue(i));
                Sub(yLocal,yLocal,x3Local,1);
                yGm.SetValue(i, yLocal.GetValue(0));
            }
            inQueueX3.FreeTensor(x3Local);
            outQueueY.FreeTensor(yLocal);
        }
    }  
}

template <typename aType, typename bType, typename cType>
__aicore__ inline void
KernelMatMulSub<aType, bType, cType>::CalcOffset(int32_t blockIdx, const TCubeTiling &tiling,
                                                             int32_t &offsetA, int32_t &offsetB, int32_t &offsetC)
{
    auto mSingleBlocks = Ceiling(tiling.M, tiling.singleCoreM);
    auto mCoreIndx = blockIdx % mSingleBlocks;
    auto nCoreIndx = blockIdx / mSingleBlocks;

    offsetA = mCoreIndx * tiling.Ka * tiling.singleCoreM;
    offsetB = nCoreIndx * tiling.singleCoreN;
    offsetC = mCoreIndx * tiling.N * tiling.singleCoreM + nCoreIndx * tiling.singleCoreN;
    // printf("CalcOffset:%d %d %d\r\n ",mSingleBlocks,mCoreIndx,nCoreIndx );

    tailM = tiling.M - mCoreIndx * tiling.singleCoreM;
    tailM = tailM < tiling.singleCoreM ? tailM : tiling.singleCoreM;
    tailN = tiling.N - nCoreIndx * tiling.singleCoreN;
    tailN = tailN < tiling.singleCoreN ? tailN : tiling.singleCoreN;
    if (tailM < tiling.singleCoreM || tailN < tiling.singleCoreN) {
        matmulObj.SetTail(tailM, tailN);
    }
    // printf("SetTail:%d %d\r\n ",tailM,tailN );
}


#endif
